import torch
import torch.nn as nn

import matplotlib.pyplot as plt
import d2l.torch as d2l


class PositionalEncoding(nn.Module):
    """位置编码"""
    def __init__(self, num_hiddens, dropout, max_len=1000):
        super(PositionalEncoding, self).__init__()
        # 定义 Dropout 层，防止过拟合
        self.dropout = nn.Dropout(dropout)
        # 创建一个足够长的位置编码矩阵 P
        self.P = torch.zeros((1, max_len, num_hiddens))
        # 计算用于生成位置编码的 X
        X = torch.arange(max_len, dtype=torch.float32).reshape(
            -1, 1) / torch.pow(10000, torch.arange(
            0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)
        # 将正弦函数结果赋值给 P 的偶数位置
        self.P[:, :, 0::2] = torch.sin(X)
        # 将余弦函数结果赋值给 P 的奇数位置
        self.P[:, :, 1::2] = torch.cos(X)
 
    def forward(self, X):
        # 将输入 X 与对应长度的位置编码相加
        X = X + self.P[:, :X.shape[1], :].to(X.device)
        # 对加上位置编码后的 X 应用 Dropout 操作并返回
        return self.dropout(X)


encoding_dim, num_steps = 32, 60
pos_encoding = PositionalEncoding(encoding_dim, 0)
pos_encoding.eval()
X = pos_encoding(torch.zeros((1, num_steps, encoding_dim)))
P = pos_encoding.P[:, :X.shape[1], :]

plt.figure(figsize=(6, 2.5))
for d in range(8):
    plt.plot(torch.arange(num_steps), P[0, :, d], label=f'Col {d}')
plt.legend()
plt.xlabel('Row (position)')
plt.show()